package com.arthas.client;

import org.apache.commons.net.telnet.TelnetClient;
import org.apache.commons.net.telnet.WindowSizeOptionHandler;

import java.io.IOException;
import java.io.InputStream;
import java.io.InputStreamReader;
import java.io.OutputStream;

public class TelnetFramework {
    private InputStreamReader responseReader;
    private TelnetClient telnetClient;

    private long timeout = 5000L;

    /**
     * Connect to a device using telnet.
     *
     * @param telnetHost the telnet host IP address
     * @param telnetPort the telnet port
     * @throws IOException
     */
    public synchronized void connect(String telnetHost, int telnetPort) throws IOException {
        if ((telnetClient == null) || ((telnetClient != null) && (!telnetClient.isConnected()))) {
            telnetClient = new TelnetClient(telnetHost);
            telnetClient.connect(telnetHost, telnetPort);
        }
        responseReader = new InputStreamReader(telnetClient.getInputStream());
    }

    /**
     * Disconnect a telnet connection to a device
     *
     * @throws IOException
     */
    public synchronized void disconnect() throws IOException {
        if ((telnetClient != null) && (telnetClient.isConnected())) {
            responseReader.close();
            telnetClient.disconnect();
        }
    }

    /**
     * @return
     */
    public synchronized String write(String command, String[] waitFor) throws IOException {
        telnetClient.getOutputStream().write(command.getBytes());
        telnetClient.getOutputStream().flush();
        if (waitFor != null) {
            return waitFor(waitFor);
        }
        return null;
    }

    /**
     * Tests if the telnet client instance is connected
     *
     * @return true if it is connected; false otherwise
     */
    public boolean isConnected() {
        boolean connected = false;
        if (telnetClient != null) {
            connected = telnetClient.isConnected();
        }
        return connected;
    }

    /**
     *
     */
    public String waitFor(String[] waitForArray) throws IOException {
        StringBuffer answerFromRemoteHost = new StringBuffer();
        boolean found = false;
        do {
            char readChar = 0;
            long currentTime = System.currentTimeMillis();
            long timeoutTime = currentTime + timeout;
            while (readChar == 0) {
                if (responseReader == null) {
                    // responseReader can only be set to null if method
                    // releaseTelnetInputStreamReader()
                    // has been called, which should happen if host becomes
                    // unavailable.
                    throw new IOException("Telnet host is unavailable; stopped waiting for answer.");
                }

                if (responseReader.ready()) {
                    readChar = (char) responseReader.read();
                } else {
                    try {
                        Thread.sleep(50);
                    } catch (InterruptedException e) {
                        // Do nothing
                    }
                }

                currentTime = System.currentTimeMillis();
                if ((!responseReader.ready()) && (currentTime > timeoutTime)) {
                    throw new IOException("A timeout has occured when trying to read the telnet stream");
                }
            }

            answerFromRemoteHost.append(readChar);

            for (String aWaitFor : waitForArray) {
                found = answerFromRemoteHost.toString().contains(aWaitFor);
            }

        } while (!found);
        return answerFromRemoteHost.toString();
    }

    /**
     * Retrieves the input stream associated to this telnet connection
     */
    public InputStream getInputStream() {
        InputStream s = null;
        if (telnetClient != null) {
            s = telnetClient.getInputStream();
        }
        return s;
    }

    /**
     * Retrieves the output stream associated to this telnet connection
     */
    public OutputStream getOutputStream() {
        OutputStream s = null;
        if (telnetClient != null) {
            s = telnetClient.getOutputStream();
        }
        return s;
    }

    /**
     * 调整客户端窗口大小
     *
     * @param nWidth  每行字符串长度,包括命令提示符
     * @param nHeight
     */
    public void setWindowSize(int nWidth, int nHeight) {
        WindowSizeOptionHandler sizeOptionHandler = new WindowSizeOptionHandler(nWidth, nHeight);
        try {
            telnetClient.sendSubnegotiation(sizeOptionHandler.startSubnegotiationLocal());
        } catch (IOException e) {
            e.printStackTrace();
        }
    }

    public long getTimeout() {
        return timeout;
    }

    public void setTimeout(long timeout) {
        this.timeout = timeout;
    }
}
